#!/usr/bin/python
import os
import sys
import time
import json
import glob
import logging
import getpass
import datetime
import ConfigParser
import subprocess
from multiprocessing import Process, Manager, Pool

import yaml
import psutil

ABS_PATH = os.path.abspath(__file__)
BASE_DIR = os.path.dirname(ABS_PATH)
BASE_NAME = os.path.basename(ABS_PATH)
LOG_DIR = os.path.join(BASE_DIR, 'log')
CONF_DIR = os.path.join(BASE_DIR, 'conf')


logging.basicConfig(filename=os.path.join(LOG_DIR, 'fullZone.log'), format='%(asctime)s %(message)s', level=logging.INFO)

cf = ConfigParser.ConfigParser()
with open(os.path.join(CONF_DIR, 'conf.ini'), 'r') as f:
    cf.readfp(f)

zcounter_exec_path = cf.get('zcounter', 'exec_path')
zcounter_pid_path = cf.get('zcounter', 'pid_path')
zcounter_dat_path = cf.get('zcounter', 'dat')

cnzonefull_exec_path = cf.get('cnzonefull', 'exec_path')
cnzonefull_pid_path = cf.get('cnzonefull', 'pid_path')
cnzonefull_dat_path = cf.get('cnzonefull', 'dat_path')

idnzonefull_exec_path = cf.get('idnzonefull', 'exec_path')
idnzonefull_pid_path = cf.get('idnzonefull', 'pid_path')
idnzonefull_dat_path = cf.get('idnzonefull', 'dat_path')

zone_soa_path = cf.get('other', 'soa_path')
zone_list_file = cf.get('other', 'zone_list_file')

domain_dic_path = cf.get('vipdomain', 'domain_dic_path')
ns_dic_path = cf.get('vipdomain', 'ns_dic_path')
tld_dic_path = cf.get('vipdomain', 'tld_dic_path')

PROCESS_NUM = int(cf.get('program', 'process_num'))
SPLIT_LINE = int(cf.get('program', 'split_line'))

dig = cf.get('other', 'dig_path')
err_dat_path = cf.get('other', 'err_dat_path')
tmp_dat_path = cf.get('other', 'tmp_dat_path')
named_checkzone_bin = cf.get('other', 'named_checkzone_bin')

prefix_dat_path = cf.get('prefix', 'prefix_dat_path')
soa_template_path = cf.get('prefix', 'soa_template_path')
glue_dat_path = cf.get('prefix', 'glue_dat_path')
added_soa_num = cf.get('prefix', 'added_soa_num')


prefix_dat_path = cf.get('prefix', 'prefix_dat_path')
cn_dat_path = cf.get('cnzonefull', 'dat_path')
idn_dat_path = cf.get('idnzonefull', 'dat_path')
final_zone_path = cf.get('prefix', 'final_zone_path')
checkzone_log_path = cf.get('other', 'named_checkzone_log_path')


def setenv():
    pass


def check_user():
    user = getpass.getuser()
    if user != "ops":
        message = '{0} is not allowed to run {1}.'.format(user, BASE_NAME)
        logging_and_print_message(message, 'error')
        sys.exit(1)


def check_pid():
    process_lst = [psutil.Process(i) for i in psutil.pids()]
    lst = filter(lambda x: BASE_NAME in x.cmdline() and 'python' in x.cmdline(), process_lst)
    if len(lst) > 1:
        message = '{0} is already running! Please check!'.format(ABS_PATH)
        logging_and_print_message(message, 'error')
        sys.exit(1)


def logging_and_print_message(message, message_type):
    print(message)
    if message_type == 'info':
        logging.info(message)
    if message_type == 'error':
        logging.error(message)


def sys_cmd(cmd):
#   print(cmd)
    sp = subprocess.Popen(cmd, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out, err = sp.communicate()
    return out, err


def sys_script(script_path, pid_path, process_name, dic):
    message = 'Start running {0}'.format(script_path)
    logging_and_print_message(message, 'info')

    out, err = sys_cmd(script_path)

    with open(pid_path, 'r') as f:
        pid = f.read()
    DATA_GENERATED_THIS_TIME = False
    while True:
        if psutil.pid_exists(int(pid)):
            DATA_GENERATED_THIS_TIME = True
            time.sleep(5) 
        else:
            break

    if DATA_GENERATED_THIS_TIME: 
        message = 'Finished {0}'.format(script_path)
        logging_and_print_message(message, 'info')
        dic[process_name] = True
    else:
        message = '{0} runs abnormally this time, please check.'.format(script_path)
        logging_and_print_message(message, 'error')
        dic[process_name] = False


def generate_assume_data(name, dic):
    sys_script(zcounter_exec_path, zcounter_pid_path, name, dic)


def generate_cn_zone(name, dic):
    sys_script(cnzonefull_exec_path, cnzonefull_pid_path, name, dic)


def generate_idn_zone(name, dic):
    sys_script(idnzonefull_exec_path, idnzonefull_pid_path, name, dic)


def generate_data(process_dic):
    process_lst = [Process(target=generate_assume_data, args=('assume_data', process_dic)),
                   Process(target=generate_cn_zone, args=('cn_zone', process_dic)),
                   Process(target=generate_idn_zone, args=('idn_zone', process_dic))]
    
    for i in process_lst:
        i.start()
    for i in process_lst:
        i.join()


def get_assume_data():
    try:
        with open(zcounter_dat_path, 'r') as f:
            lines = f.readlines()
            if lines[-1].find('EOF') > 0:
                assume_dic = {}
                for line in lines[:-2]:
                    if line:
                        zone, num = line[:-1].split('\t')
                        assume_dic.update({zone: num})
                message = 'Get assume data from {0}'.format(zcounter_dat_path)
                logging_and_print_message(message, 'info')
                return assume_dic
            else:
                message = 'Not found ";EOF" in {0}.'.format(zcounter_dat_path)
                logging_and_print_message(message, 'error')
                sys.exit(1)
    except Exception as e:
        logging.error(e)
        print(e)
        sys.exit(1)


def check_zone_end_tag(path, dic={}):
    os.chdir(path)
    for zone in filter(lambda x: x.endswith('zone'), os.listdir(path)):
        with open(zone, 'r') as f:
            lines = f.readlines()
        if ';EOF' in lines[-1]:
            dic[zone] = {'EOF': True, 'count': len(lines) - 2}
        else:
            dic[zone] = {'EOF': False, 'count': len(lines) - 2}
    return dic

def check_cn_zone_completion():
    return check_zone_end_tag(cnzonefull_dat_path, {})

def check_idn_zone_completion():
    return check_zone_end_tag(idnzonefull_dat_path, {})

def check_zone_completion():
    check_info = {'cn': check_cn_zone_completion(), 'idn': check_idn_zone_completion()}
    error_info = {}
    for t, zones_dic in check_info.items():
        for zone, zone_dic in zones_dic.items():
            if not zone_dic['EOF']:
                error_info.setdefault(t, {}).update({zone: zone_dic})
    if error_info:
        for t, zones_dic in error_info.items():
            message = 'Failed checking {0}_zone_completion, uncompleted zones: {1}.'.format(t, zones_dic.keys())
            logging_and_print_message(message, 'error')
        sys.exit(1)
    else:
        message = 'Finished check_zone_completion.'
        logging_and_print_message(message, 'info')
        return check_info

def check_assume_data_and_zone_data(assume_data, zone_data):
    pass


def check_fast_update():
    pass



def check_if_domain_is_vip_domain(domain, ns, vip_domain_lst, dic):
    if domain in vip_domain_lst:
        dic.setdefault(domain, set())
        s = dic[domain]
        s.add(ns)
        dic[domain] = s


def check_if_ns_is_vip_ns(ns, qtype, glue, vip_ns_lst, dic):
    if ns in vip_ns_lst:
        dic.setdefault(ns, {})
        ns_dic = dic[ns]
        ns_dic.setdefault(qtype, set()).add(glue)
        dic[ns] = ns_dic


def filter_zone_data_vip_domain_and_vip_ns(n, lines, vip_domain_dic, vip_ns_dic, domain_dic, ns_dic):
    for line in lines:
        # exclude the last two line
        if len(line.split()) == 0 or line.startswith(';'):
            continue
        try:
            if line.split()[2] == 'NS':
                domain, qclass, qtype, ns = [x.lower() for x in line.split()]
                check_if_domain_is_vip_domain(domain, ns, vip_domain_dic.keys(), domain_dic)
                continue
            if line.split()[2] in ['A', 'AAAA']:
                ns, qclass, qtype, glue = [x.lower() for x in line.split()]
                check_if_ns_is_vip_ns(ns, qtype, glue, vip_ns_dic.keys(), ns_dic)
                continue
        except Exception as e:
            logging_and_print_message(e, 'error')

    logging_and_print_message('Finished part {0}.'.format(n), 'info')


def get_vip_domain_and_vip_ns(zonefull_dat_path, vip_domain_dic, vip_ns_dic, domain_dic, ns_dic):
    logging_and_print_message('Begin to load zones at {0}.'.format(zonefull_dat_path), 'info')
    pool = Pool(processes=PROCESS_NUM)
    n = 0
    lines = []
    for zone in glob.glob(zonefull_dat_path+'/*zone'):
        zone_data_path = os.path.join(zonefull_dat_path, zone)
        with open(zone_data_path, 'r') as f:
            lines.extend(f.readlines()) 
            logging_and_print_message('Got {0} lines from zone {1}.'.format(len(lines), zone), 'info')
            while len(lines) > SPLIT_LINE:
                n += 1
                logging_and_print_message('Start part {0}.'.format(n), 'info')
                pool.apply_async(filter_zone_data_vip_domain_and_vip_ns, (n, lines[:SPLIT_LINE], vip_domain_dic, vip_ns_dic, domain_dic, ns_dic))
                lines = lines[SPLIT_LINE:]
    # the rest of lines 
    n += 1
    logging_and_print_message('Start part {0}.'.format(n), 'info')
    pool.apply_async(filter_zone_data_vip_domain_and_vip_ns, (n, lines, vip_domain_dic, vip_ns_dic, domain_dic, ns_dic))

    pool.close()
    pool.join()


def filter_diff_domain_and_ns(vip_domain_dic, vip_ns_dic, domain_dic, ns_dic):
    diff_domain_dic = {}
    for vip_domain, vip_ns_lst in vip_domain_dic.items():
        ns_lst = domain_dic.get(vip_domain, None)
        if ns_lst and len(set(vip_ns_lst) - set(ns_lst)) == 0 and len(set(ns_lst) - set(vip_ns_lst)) == 0:
            continue
        diff_domain_dic[vip_domain] = {'from_vipdomain': vip_ns_lst, 'from_zone_data': ns_lst}

    diff_ns_dic = {}
    for vip_ns, vip_qtype_dic in vip_ns_dic.items():
        qtype_dic = ns_dic.get(vip_ns, None)
        # qtype_dic is None OR ns not has same type of qtype: A, AAAA
        if not qtype_dic or len(set(vip_qtype_dic.keys()) - set(qtype_dic.keys())) != 0 or len(set(qtype_dic.keys()) - set(vip_qtype_dic.keys())) != 0:
            diff_ns_dic[vip_ns] = {'from_vipdomain': vip_qtype_dic, 'from_zone_data': qtype_dic}
            continue
        # ns has the same type of qtype
        for vip_qtype, vip_glue_lst in vip_qtype_dic.items():
            # if vip_ns has A and AAAA, but from zone data only get A
            glue_lst = qtype_dic.get(vip_qtype, None)
            if not glue_lst or len(set(vip_glue_lst) - set(glue_lst)) != 0 or len(set(glue_lst) - set(vip_glue_lst)) != 0:
                diff_ns_dic[vip_ns] = {'from_vipdomain': vip_qtype_dic, 'from_zone_data': qtype_dic}
                break

    return diff_domain_dic, diff_ns_dic


def check_vipdomain():
    # vip_domain_dic: {u'cass.cn.': ['ns1.cass.cn.', 'ns2.cass.cn.']}
    with open(domain_dic_path, 'r') as f:
        vip_domain_dic = json.load(f)
    # vip_ns_dic: {'ns2.cass.cn.': {'A': ['103.247.176.2']}, 'ns1.cass.cn.': {'A': ['103.247.176.1']}}
    with open(ns_dic_path, 'r') as f:
        vip_ns_dic = json.load(f)
    
    domain_dic = Manager().dict()
    ns_dic = Manager().dict()

    logging_and_print_message('Process num: {0}, split lines: {1}.'.format(PROCESS_NUM, SPLIT_LINE), 'info')
    for path in [cnzonefull_dat_path, idnzonefull_dat_path ]:
        get_vip_domain_and_vip_ns(path, vip_domain_dic, vip_ns_dic, domain_dic, ns_dic)

    domain_dic = domain_dic.copy()
    ns_dic = ns_dic.copy()
    for domain, ns_set in domain_dic.items():
        domain_dic[domain] = list(ns_set)
    for ns, qtype_dic in ns_dic.items():
        for qtype, glue_set in qtype_dic.items():
            qtype_dic[qtype] = list(glue_set)

    diff_domain_dic, diff_ns_dic = filter_diff_domain_and_ns(vip_domain_dic, vip_ns_dic, domain_dic, ns_dic)

    with open(os.path.join(tmp_dat_path, 'domain_dic'), 'w') as f:
        json.dump(domain_dic, f, indent=4)
        message = 'Saved tmp file: {0} at {1}.'.format('domain_dic', tmp_dat_path)
        logging_and_print_message(message, 'info')

    with open(os.path.join(tmp_dat_path, 'ns_dic'), 'w') as f:
        json.dump(ns_dic, f, indent=4)
        message = 'Saved tmp file: {0} at {1}.'.format('ns_dic', tmp_dat_path)
        logging_and_print_message(message, 'info')

    logging_and_print_message('Found {0} domains at zone data.'.format(len(domain_dic)), 'info')
    logging_and_print_message('Found {0} vip domains at vip_domain file.'.format(len(vip_domain_dic)), 'info')
    logging_and_print_message('Found {0} ns at zone data.'.format(len(ns_dic)), 'info')
    logging_and_print_message('Found {0} vip ns at vip_ns file.'.format(len(vip_ns_dic)), 'info')
    
    if len(domain_dic) - len(vip_domain_dic) != 0 or len(ns_dic) - len(vip_ns_dic) != 0:
        with open(os.path.join(err_dat_path, 'diff_vip_domain'), 'w') as f:
            json.dump(diff_domain_dic, f, indent=4)
            message = 'Saved err file: {0} at {1}.'.format('diff_vip_domain', err_dat_path)
            logging_and_print_message(message, 'info')

        with open(os.path.join(err_dat_path, 'diff_vip_ns'), 'w') as f:
            json.dump(diff_ns_dic, f, indent=4)
            message = 'Saved err file: {0} at {1}.'.format('diff_vip_ns', err_dat_path)
            logging_and_print_message(message, 'info')


def get_zone_soa_from_mul_server(zone, ips):
    max_soa = 0
    for ip in ips:
        cmd = '{0} @{1} {2} soa +short'.format(dig, ip, zone)
        out, err = sys_cmd(cmd)
        if out:
            if len(out.split('\n')) == 2:
                soa = out.split()[2]
            else:
                continue
            max_soa = max(int(soa), int(max_soa))
    logging_and_print_message('Zone: {0}, SOA: {1}'.format(zone, max_soa), 'info')
    return max_soa


def get_zone_soa():
    soa_dic = {}
    with open(glue_dat_path, 'r') as f:
        dic = yaml.load(f)
    ips = [ip for _, ip in cf.items('servers')]
    logging_and_print_message('Get soa for each zone from servers: {0}.'.format(ips), 'info')
    for zone in dic['zone_info'].keys():
        soa = get_zone_soa_from_mul_server(zone, ips)
        soa_dic[zone] = soa

    return soa_dic



def generate_prefix(soa_dic):
    with open(soa_template_path, 'r') as f:
        lines = f.readlines()

    zone_ns = set()
    ns_glue = set()

    with open(glue_dat_path, 'r') as f:
        dic = yaml.load(f)
    for zone, soa in soa_dic.items():
    #   print(zone, soa)
        with open(os.path.join(prefix_dat_path, zone), 'w') as f:
            for line in lines:
                if line.find("REPLACE_SOA") > 0:
                    line = line.replace("REPLACE_SOA", str(int(soa)+int(added_soa_num)))
                f.write(line)

            for ns_cop in dic['zone_info'][zone]:
                for ns in ns_cop:
                    f.write('{0}\tIN\tNS\t{1}\n'.format(zone, ns))
                    # exclude the xn--fiqs8s & xn--fiqz9s
                    if zone.endswith('cn.'):
                        zone_ns.add((zone, ns))
            for ns_cop in dic['zone_info'][zone]:
                for ns in ns_cop:
                    qtype_dic = dic['ns_info'].get(ns)
                    if qtype_dic:
                        for qtype, glue_lst in qtype_dic.items():
                            for glue in glue_lst:
                              # f.write('{0}\tIN\t{1}\t{2}\n'.format(ns, qtype, glue))
                                # only like "mil.cn IN NS ns3.mil.cn" need glue
                                if zone in ns:
                                    f.write('{0}\tIN\t{1}\t{2}\n'.format(ns, qtype, glue))
                                # exclude the xn--fiqs8s & xn--fiqz9s
                                if zone.endswith('cn.'):
                                    ns_glue.add((ns, qtype, glue))
        # write to cn. second time
        with open(os.path.join(prefix_dat_path, 'cn.'), 'w') as f:
            for line in lines:
                if line.find("REPLACE_SOA") > 0:
                    line = line.replace("REPLACE_SOA", str(int(soa)+int(added_soa_num)))
                f.write(line)
            for zone, ns in sorted(zone_ns):
                f.write('{0}\tIN\tNS\t{1}\n'.format(zone, ns))
            for ns, qtype, glue in sorted(ns_glue):
                f.write('{0}\tIN\t{1}\t{2}\n'.format(ns, qtype, glue))


def generate_sld_zone_file():
    os.chdir(cn_dat_path)
    zones = filter(lambda x: x.endswith('zone'), os.listdir(os.getcwd()))
    for zone in zones:
        sld = zone.split('.')[0]
        if sld != 'cn':
            prefix_file = os.path.join(prefix_dat_path, sld+'.cn.')
            with open(prefix_file, 'r') as f:
                lines = f.readlines()
            final_zone_file = os.path.join(final_zone_path, sld+'.cn.zone')
            with open(final_zone_file, 'w') as f:
                f.writelines(lines)
                with open(os.path.join(cn_dat_path, zone), 'r') as fp:
                    f.writelines(fp.readlines())
            logging_and_print_message('Finished generating {0}.'.format(final_zone_file), 'info')

def generate_idn_zone_file():
    os.chdir(idn_dat_path)
    zones = filter(lambda x: x.endswith('zone'), os.listdir(os.getcwd()))
    for zone in zones:
        tld = zone.split('.')[0]
        if tld != 'cn':
            prefix_file = os.path.join(prefix_dat_path, tld+'.')
            with open(prefix_file, 'r') as f:
                lines = f.readlines()
            final_zone_file = os.path.join(final_zone_path, zone)
            with open(final_zone_file, 'w') as f:
                f.writelines(lines)
                with open(os.path.join(idn_dat_path, zone), 'r') as fp:
                    f.writelines(fp.readlines())
            logging_and_print_message('Finished generating {0}.'.format(final_zone_file), 'info')

def generate_cn_zone_file():
    with open(os.path.join(prefix_dat_path, 'cn.'), 'r') as f:
        lines = f.readlines()
    final_zone_file = os.path.join(final_zone_path, 'cn.zone')
    with open(final_zone_file, 'w') as f:
        f.writelines(lines)
        with open(os.path.join(cn_dat_path, 'cn.zone'), 'r') as fp:
            f.writelines(fp.readlines())
        with open(os.path.join(idn_dat_path, 'cn.zone'), 'r') as fp:
            f.writelines(fp.readlines())
        logging_and_print_message('Finished generating {0}.'.format(final_zone_file), 'info')
        

def generate_zone_file():
    logging_and_print_message('Begin to generate final zone.', 'info')
    generate_sld_zone_file()
    generate_idn_zone_file()
    generate_cn_zone_file()



def async_checkzone(zone, final_zone_file, dic):
    logging_and_print_message('Begin to check {0} at {1}.'.format(zone, final_zone_file), 'info')
    cmd = '{0} -i local {1} {2}'.format(named_checkzone_bin, zone, final_zone_file)
    out, err = sys_cmd(cmd)
    with open(os.path.join(checkzone_log_path, zone), 'w') as f:
        f.writelines(out)

    if out.split('\n')[-2] == 'OK':
        dic[zone] = True
    else:
        dic[zone] = False

    logging_and_print_message('Finished check {0}'.format(zone), 'info')

def named_checkzone():
    logging_and_print_message('Begin to named-checkzone.', 'info')
    pool = Pool(processes=PROCESS_NUM)
    os.chdir(final_zone_path)
    dic = Manager().dict()
    for zone_file in os.listdir(final_zone_path):
        zone = '.'.join(zone_file.split('.')[:-1])
        final_zone_file = os.path.join(final_zone_path, zone_file)
        pool.apply_async(async_checkzone, (zone, final_zone_file, dic))

    pool.close()
    pool.join()
    logging_and_print_message('Finished named-checkzone.', 'info')
#   print(json.dumps(dic.copy(), indent=4))
    if all(dic.values()):
        logging_and_print_message('All zone passed checked.', 'info')
    else:
        lst = []
        for zone, v in dic.items():
            if not v:
                lst.append(zone)
        logging_and_print_message('Failed named-checkzone zones: {0}.'.format(lst), 'error')


def main():
#   setenv()
#   check_user()
#   check_pid()
#   check_fast_update()
#   process_dic = Manager().dict()
#   generate_data(process_dic)
#   print(json.dumps(process_dic.copy(), indent=4))
#   if False in process_dic.values():
#       logging_and_print_message('Error occured in generate data.', 'info')
#       sys.exit(1)

#   assume_data = get_assume_data()
#   print(json.dumps(assume_data, indent=4))

#   zone_data = check_zone_completion()
#   print(json.dumps(zone_data, indent=4))

#   check_assume_data_and_zone_data(assume_data, zone_data)

#   check_vipdomain()
#   soa_dic = get_zone_soa()
#   generate_prefix(soa_dic)
#   generate_zone_file()
    named_checkzone()

    
if __name__ == '__main__':
    main()
#   test()
